/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <atomic>
#include <climits>
#include <condition_variable>
#include <inttypes.h>
#include <mutex>
#include <queue>
#include <string>
#include <map>
#include <vector>
#include <algorithm>
#include <sstream>
#include "utils_log.h"

namespace El {
namespace Utils {

#define TIME_MAX_THRESHOLD 200

class TimeCost {
public:
    explicit TimeCost(const std::string &name = "TimeCost", bool mark_start = false)
        : name_(name), threshold_(TIME_MAX_THRESHOLD)
    {
        if (mark_start) {
            Mark("TimeCost Start");
        }
    }

    ~TimeCost() {}

    void SetThreshold(uint64_t threshold)
    {
        threshold_ = threshold;
    }

    // 记录一个新的时间点
    void Mark(const std::string &point_name)
    {
        uint64_t current_time =
            std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now().time_since_epoch())
                .count();
        timePoints_.push_back({point_name, current_time});
    }

    // 打印统计时间，check_threshold为true时只打印超过阈值的记录
    void Print(bool check_threshold = false)
    {
        if (timePoints_.size() < 2) {
            return;
        }

        bool has_content_to_print = false;
        std::stringstream ss;

        // 如果需要检查阈值，先检查是否有超过阈值的记录
        if (check_threshold) {
            for (size_t i = 0; i < timePoints_.size() - 1; ++i) {
                uint64_t duration = timePoints_[i + 1].second - timePoints_[i].second;
                if (duration >= threshold_) {
                    has_content_to_print = true;
                    break;
                }
            }

            // 如果没有超过阈值的记录，直接返回
            if (!has_content_to_print) {
                return;
            }
        }

        ss << "\n=== " << name_ << " TimeCost Statistics ===\n";

        // 按记录顺序打印每个阶段的耗时
        for (size_t i = 0; i < timePoints_.size() - 1; ++i) {
            const auto &start = timePoints_[i];
            const auto &end = timePoints_[i + 1];
            uint64_t duration = end.second - start.second;

            // 如果需要检查阈值且未超过阈值，则跳过
            if (check_threshold && duration < threshold_) {
                continue;
            }

            ss << start.first << " -> " << end.first << ": " << duration << " ms";
            if (duration >= threshold_) {
                ss << " [WARNING: Exceeds threshold]";
            }
            ss << "\n";
        }

        // 打印总耗时
        uint64_t total = timePoints_.back().second - timePoints_.front().second;
        ss << "Total time (" << timePoints_.front().first << " -> " << timePoints_.back().first << "): " << total
           << " ms\n";
        ss << "===========================\n";

        EL_INFO("{}", ss.str().c_str());
        timePoints_.clear(); // 打印后自动清除时间点
    }

    // 清除所有记录的时间点
    void Clear()
    {
        timePoints_.clear();
    }

private:
    std::vector<std::pair<std::string, uint64_t>> timePoints_; // 按记录顺序存储时间点
    const std::string name_;                                   // 统计对象名称
    uint64_t threshold_;                                       // 耗时阈值
};

} // namespace Utils
} // namespace El